#!/bin/bash

# 设置TensorRT路径
export TENSORRT_ROOT=/data2/xd/TensorRT-8.5.1.7.Linux.x86_64-gnu.cuda-11.8.cudnn8.6/TensorRT-8.5.1.7
export PATH="$PATH:$TENSORRT_ROOT/bin"
export LD_LIBRARY_PATH="$TENSORRT_ROOT/lib:$LD_LIBRARY_PATH"

# 设置插件路径
PLUGIN_PATH="/data2/xd/Co-DETR-TensorRT/co_detr_ros_ws/src/co_detr_ros/grid_sampler/build/libtrtgrid_sampler.so"

# 检查插件文件是否存在
if [ ! -f "$PLUGIN_PATH" ]; then
    echo "错误: 插件文件 $PLUGIN_PATH 不存在"
    exit 1
fi

echo "=== 优化Co-DETR TensorRT引擎构建 ==="
echo "插件路径: $PLUGIN_PATH"

# 构建优化的FP16引擎
echo "构建优化的FP16引擎..."
trtexec \
    --onnx=models/co_detr_sim.onnx \
    --saveEngine=models/co_detr_optimized_fp16.plan \
    --memPoolSize=workspace:16000 \
    --fp16 \
    --plugins="$PLUGIN_PATH" \
    --tacticSources=+CUBLAS,+CUBLAS_LT,+CUDNN \
    --noBuilderCache \
    --verbose

if [ $? -eq 0 ]; then
    echo "✅ 优化的FP16引擎构建成功!"
else
    echo "❌ FP16引擎构建失败!"
    exit 1
fi

# 构建优化的FP32引擎
echo "构建优化的FP32引擎..."
trtexec \
    --onnx=models/co_detr_sim.onnx \
    --saveEngine=models/co_detr_optimized_fp32.plan \
    --memPoolSize=workspace:16000 \
    --tacticSources=+CUBLAS,+CUBLAS_LT,+CUDNN \
    --noBuilderCache \
    --verbose

if [ $? -eq 0 ]; then
    echo "✅ 优化的FP32引擎构建成功!"
else
    echo "❌ FP32引擎构建失败!"
    exit 1
fi

# 构建混合精度引擎（推荐）
echo "构建混合精度引擎（推荐）..."
trtexec \
    --onnx=models/co_detr_sim.onnx \
    --saveEngine=models/co_detr_mixed_precision.plan \
    --memPoolSize=workspace:16000 \
    --fp16 \
    --int8 \
    --plugins="$PLUGIN_PATH" \
    --tacticSources=+CUBLAS,+CUBLAS_LT,+CUDNN \
    --noBuilderCache \
    --verbose

if [ $? -eq 0 ]; then
    echo "✅ 混合精度引擎构建成功!"
else
    echo "❌ 混合精度引擎构建失败!"
fi

echo "=== 引擎构建完成 ==="
echo "可用引擎文件："
echo "- co_detr_optimized_fp16.plan (优化的FP16)"
echo "- co_detr_optimized_fp32.plan (优化的FP32)"
echo "- co_detr_mixed_precision.plan (混合精度，推荐)" 